Skip to content

[Perf] Add fused qk rotary embedding kernel for Qwen2.5-VL ViT #715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 8 commits into from
Sep 30, 2025
Merged

[Perf] Add fused qk rotary embedding kernel for Qwen2.5-VL ViT #715

merged 8 commits into from
Sep 30, 2025

Conversation

kliuae-amd
Copy link

@kliuae-amd kliuae-amd commented Sep 26, 2025 •
edited by github-actions bot
Loading

Purpose

This PR targets the Qwen2.5-VL models, unifying the rotary embedding of query and key in ViT into one kernel adopted from flash_attn. The kernel speeds up Qwen2 ViT's rotary by 1.9x compared with the calculating q and k embeddings separately.

Kernel benchmark

python benchmarks/kernels/benchmark_vision_rotary_emb.py   --batch-size 2 --seq-len 16384 --num-heads 16 --head-size 80   --dtype bfloat16 --device cuda --warmup-iter 10 --benchmark-iter 1000

Config: batch=2, seqlen=16384, heads=16, head_dim=80, dtype=torch.bfloat16
Kernel mean runtime (ms) median runtime (ms)
1c, separated q and k 0.6679 0.6590
2c, fused q and k 0.3451 0.3416 
Fusion speedup: 1.936x

Test Plan

Qwen2-VL-72B-Instruct correctness validation on MathVista

Server command:

VLLM_USE_TRITON_FLASH_ATTN=0 \
VLLM_USE_V1=1 \
VLLM_ROCM_USE_AITER=1 \
SAFETENSORS_FAST_GPU=1 \
VLLM_DISABLE_COMPILE_CACHE=1 \
vllm serve Qwen/Qwen2.5-VL-72B-Instruct \
    --tensor_parallel_size 2 \
    --trust_remote_code \
    --no_enable_prefix_caching \
    --enable_multimodal_encoder_data_parallel \
    --disable-mm-preprocessor-cache

Test Result

Eval

Command:
python3 -m eval.run eval_vllm \
        --model_name Qwen/Qwen2.5-VL-72B-Instruct \
        --url http://0.0.0.0:8088 \
        --output_dir temp_mistral_eval \
        --eval_name "mathvista"

Separated qk (before):

Metrics:                                                                                                                                                                                                    
{                                                                                                                                                                                                           
    "explicit_prompt_relaxed_correctness": 0.743,                                                                                                                                                           
    "anywhere_in_answer_relaxed_correctness": 0.78                                                                                                                                                          
}

Fused qk (after):

Metrics:
{
    "explicit_prompt_relaxed_correctness": 0.749,
    "anywhere_in_answer_relaxed_correctness": 0.787
}

Benchmark

Command:
vllm bench serve  \
  --backend openai-chat  \
  --endpoint-type openai-chat  \
  --model Qwen/Qwen2.5-VL-72B-Instruct  \
  --dataset-name hf  \
  --dataset-path lmarena-ai/VisionArena-Chat  \
  --hf-split train  \
  --endpoint /v1/chat/completions  \
  --num-prompts 1000  \
  --max-concurrency 64
Metric Separated qk Fused qk
Request throughput (req/s) 1.71 1.76
Output token throughput (tok/s) 199.55 205.61
Total token throughput (tok/s) 360.85 371.16
Mean TTFT (ms) 3250.92 3556.19
Median TTFT (ms) 1749.78 1784.01
P99 TTFT (ms) 23095.19 23362.85
Mean TPOT (ms) 304.71 293.01
Median TPOT (ms) 300.05 289.61
P99 TPOT (ms) 496.55 443.96
Mean ITL (ms) 297.91 284.26
Median ITL (ms) 46.29 45.83
P99 ITL (ms) 2721.27 3096.04
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

@kliuae
add fused qk rope
2c4a086 
Signed-off-by: kliuae <[email protected]>
@kliuae
lint
569724f 
Signed-off-by: kliuae <[email protected]>
@kliuae
comment
e8d53a8 
Signed-off-by: kliuae <[email protected]>
@kliuae
add test
89cd50e 
Signed-off-by: kliuae <[email protected]>
@kliuae
tune
9e974fa 
Signed-off-by: kliuae <[email protected]>
@kliuae
typo
0910979 
Signed-off-by: kliuae <[email protected]>
@github-actions GitHub Actions
Copy link

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

@kliuae
precommit
4e0bb34 
Signed-off-by: kliuae <[email protected]>
@kliuae
yapfk
e4105e8 
Signed-off-by: kliuae <[email protected]>
@tjtanaavllm
Copy link

LGTM

@tjtanaavllm tjtanaavllm merged commit bf15fd3 into ROCm:llama_fp8_03122025 Sep 30, 2025
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@tjtanaavllm tjtanaavllm tjtanaavllm approved these changes

@Alexei-V-Ivanov-AMD Alexei-V-Ivanov-AMD Awaiting requested review from Alexei-V-Ivanov-AMD Alexei-V-Ivanov-AMD is a code owner

@shajrawi shajrawi Awaiting requested review from shajrawi shajrawi is a code owner

@gshtras gshtras Awaiting requested review from gshtras gshtras is a code owner

@maleksan85 maleksan85 Awaiting requested review from maleksan85 maleksan85 is a code owner

@sunway513 sunway513 Awaiting requested review from sunway513 sunway513 is a code owner

@hongxiayang hongxiayang Awaiting requested review from hongxiayang hongxiayang is a code owner

@charlifu charlifu Awaiting requested review from charlifu charlifu is a code owner

@mawong-amd mawong-amd Awaiting requested review from mawong-amd mawong-amd is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@kliuae-amd @tjtanaavllm @kliuae